#include "cnn.h"
#include <math.h>
#include <string.h>

// RELU激活函数
float relu(float x) {
    return x > 0 ? x : 0;
}

// 卷积操作 (简化版)
void convolution_2d(float input[28][28], float output[24][24][16], 
                    float weights[16][1][5][5], float bias[16]) {
    for (int oc = 0; oc < 16; oc++) {
        for (int i = 0; i < 24; i++) {
            for (int j = 0; j < 24; j++) {
                float sum = 0;
                for (int ki = 0; ki < 5; ki++) {
                    for (int kj = 0; kj < 5; kj++) {
                        sum += input[i + ki][j + kj] * weights[oc][0][ki][kj];
                    }
                }
                output[i][j][oc] = relu(sum + bias[oc]);
            }
        }
    }
}

// 最大池化
void max_pooling_2d(float input[24][24][16], float output[12][12][16]) {
    for (int c = 0; c < 16; c++) {
        for (int i = 0; i < 12; i++) {
            for (int j = 0; j < 12; j++) {
                float max_val = input[i*2][j*2][c];
                max_val = fmax(max_val, input[i*2+1][j*2][c]);
                max_val = fmax(max_val, input[i*2][j*2+1][c]);
                max_val = fmax(max_val, input[i*2+1][j*2+1][c]);
                output[i][j][c] = max_val;
            }
        }
    }
}

// 全连接层
void fully_connected(float input[12][12][16], float output[10], 
                    float weights[10][12*12*16], float bias[10]) {
    float flattened[12*12*16];
    int idx = 0;
    
    // 展平输入
    for (int c = 0; c < 16; c++) {
        for (int i = 0; i < 12; i++) {
            for (int j = 0; j < 12; j++) {
                flattened[idx++] = input[i][j][c];
            }
        }
    }
    
    // 全连接计算
    for (int i = 0; i < 10; i++) {
        float sum = bias[i];
        for (int j = 0; j < 12*12*16; j++) {
            sum += flattened[j] * weights[i][j];
        }
        output[i] = sum;
    }
}

// softmax激活函数
void softmax(float input[10], float output[10]) {
    float max_val = input[0];
    for (int i = 1; i < 10; i++) {
        if (input[i] > max_val) {
            max_val = input[i];
        }
    }
    
    float sum = 0;
    for (int i = 0; i < 10; i++) {
        output[i] = expf(input[i] - max_val);
        sum += output[i];
    }
    
    for (int i = 0; i < 10; i++) {
        output[i] /= sum;
    }
}

// 完整的推理函数
int mnist_inference(uint8_t *image, float *conv1_w, float *conv1_b, 
                   float *fc1_w, float *fc1_b) {
    // 将图像数据转换为浮点数并归一化
    float input[28][28];
    for (int i = 0; i < 28; i++) {
        for (int j = 0; j < 28; j++) {
            input[i][j] = image[i * 28 + j] / 255.0f;
        }
    }
    
    // 第一层卷积
    float conv1_output[24][24][16];
    convolution_2d(input, conv1_output, 
                  (float (*)[1][5][5])conv1_w, 
                  conv1_b);
    
    // 最大池化
    float pool1_output[12][12][16];
    max_pooling_2d(conv1_output, pool1_output);
    
    // 全连接层
    float fc1_output[10];
    fully_connected(pool1_output, fc1_output, 
                   (float (*)[12*12*16])fc1_w, 
                   fc1_b);
    
    // Softmax
    float softmax_output[10];
    softmax(fc1_output, softmax_output);
    
    // 找到最大概率的索引
    int max_idx = 0;
    float max_prob = softmax_output[0];
    for (int i = 1; i < 10; i++) {
        if (softmax_output[i] > max_prob) {
            max_prob = softmax_output[i];
            max_idx = i;
        }
    }
    
    return max_idx;
}
